import time

import matplotlib.pyplot as plt
import numpy as np
from xray_scan import XRayScan
from xray_scan_collection import XRayScanCollection


class RectangleData:
    """
    用于存储矩形区域内的点，以便进行插值或其他操作。
    """
    def __init__(self, points):
        """
        初始化RectangleData对象。

        :param points: 包含(x, y, value)的点的列表，value是复数
        """
        # 对点进行排序，以便后续操作（例如按x和y坐标排序）
        # self.points = sorted(points, key=lambda p: (p[0], p[1]))
        self.points = points
        self.length = 0

    def set_length(self,length):
        self.length = length

    def get_points(self):
        """
        获取存储的点。

        :return: 包含(x, y, value)的点的列表
        """
        return self.points

    def __repr__(self):
        """
        返回对象的字符串表示。
        """
        return f"RectangleData with {len(self.points)} points"

    def get_length(self):
        return self.length

class XRayScanCollectionFFT:
    def __init__(self, xray_scan_collection):
        """
        初始化XRayScanCollectionFFT对象。

        :param xray_scan_collection: XRayScanCollection对象
        """
        self.xray_scan_collection = xray_scan_collection




    def compute_fft(self):
        """
        对XRayScanCollection中的所有XRayScan对象进行傅里叶变换，并将结果存储在sensor_data属性中。
        """
        for scan in self.xray_scan_collection.scans:
            scan.sensor_data = np.fft.ifftshift(scan.sensor_data)
            scan.sensor_data = np.fft.fft(scan.sensor_data)  # 保留复数值
            scan.sensor_data = np.fft.fftshift(scan.sensor_data)
            scan.origin_index = (len(scan.sensor_data) - 1) / 2

            # 1D FFT结果
            plt.figure(figsize=(5, 5))

            # # 幅度谱（线性尺度）
            # plt.plot(np.abs(scan.sensor_data))
            # plt.title('Sensor Signal')
            # plt.xlabel('Position')
            # plt.ylabel('Magnitude')
            # plt.show()
            # #
            # # # 幅度谱（线性尺度）
            # test = np.fft.ifft(scan.sensor_data)
            # plt.plot(np.abs(test))
            # plt.title('Sensor Signal')
            # plt.xlabel('Position')
            # plt.ylabel('Magnitude')
            # plt.show()
            # time.sleep(10000)



    def convert_to_cartesian(self):
        """
        将每个XRayScan对象的极坐标数据转换为直角坐标。

        :return: 包含(x, y, value)的列表，表示每个点的直角坐标和复数值
        """
        cartesian_data = []
        for scan in self.xray_scan_collection.scans:
            angle_rad = np.radians(scan.angle_deg)  # 将角度转换为弧度
            n = len(scan.sensor_data)
            r_values = np.arange(-scan.origin_index, n - scan.origin_index)  # 以原点为中心的径向距离

            # 保留复数值
            for r, value in zip(r_values, scan.sensor_data):
                x = r * np.cos(angle_rad)
                y = r * np.sin(angle_rad)
                cartesian_data.append((x, y, value))

        return cartesian_data

    def filter_points_in_rectangle(self, aspect_ratio):
        """
        过滤并只保留位于指定长方形内的点。长方形的中心点为原点，长和宽分别平行于x轴和y轴。

        :param aspect_ratio: 长方形的长宽比（宽度 / 高度）
        :return: 一个包含 (x, y, value) 的列表，表示位于长方形内的点
        """
        # 获取直角坐标数据
        cartesian_data = self.convert_to_cartesian()

        # # 计算长方形的尺寸
        diameter = max(np.sqrt(coord[0]**2 + coord[1]**2) for coord in cartesian_data) * 2  # 取数据点的最大直径
        diameter_up = diameter * 1
        diameter_down = diameter * 0
        height_up = diameter_up / np.sqrt(1 + aspect_ratio ** 2)  # 计算长方形的高度
        width_up = aspect_ratio * height_up  # 计算长方形的宽度

        height_down = diameter_down / np.sqrt(1 + aspect_ratio ** 2)  # 计算长方形的高度
        width_down = aspect_ratio * height_down


        #
        # 过滤位于长方形内的点
        # cartesian_data = [
        #     (x, y, value) for x, y, value in cartesian_data
        #     if -width_up / 2 <= x <= width_up / 2 and -height_up / 2 <= y <= height_up / 2 and
        #        ((x < -width_down / 2 or x > width_down / 2) or (y < -height_down / 2 or y > height_down / 2))
        # ]


        # 返回过滤后的点和长方形的宽度和高度
        return cartesian_data, width_up, height_up

    def get_rectangle_data(self, aspect_ratio):
        """
        获取存储在矩形区域内的点，并以RectangleData对象的形式返回。

        :param aspect_ratio: 长方形的长宽比（宽度 / 高度）
        :return: RectangleData对象
        """



        filtered_data, width, height = self.filter_points_in_rectangle(aspect_ratio)
        result = RectangleData(filtered_data)
        return result

    def plot_cartesian(self):
        """
        在直角坐标系中绘制点，包括所有复数值（按幅度表示）。
        """
        cartesian_data = self.convert_to_cartesian()
        x_values = [coord[0] for coord in cartesian_data]
        y_values = [coord[1] for coord in cartesian_data]
        magnitudes = [np.abs(coord[2]) for coord in cartesian_data]  # 使用复数的幅度

        plt.figure(figsize=(8, 8))
        scatter = plt.scatter(x_values, y_values, c=magnitudes, cmap='viridis', marker='o')
        plt.title("Cartesian Coordinates (Including Zero Magnitudes)")
        plt.xlabel("X")
        plt.ylabel("Y")
        plt.axis("equal")
        plt.colorbar(scatter, label="Magnitude")
        plt.show()

    def plot_filtered_cartesian(self, aspect_ratio):
        """
        在直角坐标系中绘制位于指定长方形内的点，并绘制长方形的外框。

        :param aspect_ratio: 长方形的长宽比（宽度 / 高度）
        """
        filtered_data, width, height = self.filter_points_in_rectangle(aspect_ratio)
        x_values = [coord[0] for coord in filtered_data]
        y_values = [coord[1] for coord in filtered_data]
        magnitudes = [np.abs(coord[2]) for coord in filtered_data]  # 使用复数的幅度

        plt.figure(figsize=(10, 10))
        scatter = plt.scatter(x_values, y_values, c=magnitudes, cmap='viridis', marker='o')
        plt.title(f"Filtered Cartesian Coordinates (Aspect Ratio {aspect_ratio}:1)")
        plt.xlabel("X")
        plt.ylabel("Y")
        plt.axis("equal")
        plt.colorbar(scatter, label="Magnitude")

        # 绘制长方形的外框
        # rect_x = [-width / 2, width / 2, width / 2, -width / 2, -width / 2]
        # rect_y = [-height / 2, -height / 2, height / 2, height / 2, -height / 2]
        # plt.plot(rect_x, rect_y, color='red', linestyle='--', linewidth=2, label="Rectangle Boundary")
        # plt.legend()
        plt.show()

    def plot_original_cartesian(self):
            """
            在直角坐标系中绘制原始数据的点，包括原始传感器数据。
            """
            original_cartesian_data = []
            for scan in self.xray_scan_collection.scans:
                angle_rad = np.radians(scan.angle_deg)
                n = len(scan.sensor_data)
                r_values = np.arange(-scan.origin_index, n - scan.origin_index)  # 以原点为中心的径向距离

                # 将极坐标转换为直角坐标
                for r, magnitude in zip(r_values, scan.sensor_data):  # 使用原始传感器数据
                    x = r * np.cos(angle_rad)
                    y = r * np.sin(angle_rad)
                    original_cartesian_data.append((x, y, magnitude))

            # 提取x, y和magnitude值用于绘图
            x_values = [coord[0] for coord in original_cartesian_data]
            y_values = [coord[1] for coord in original_cartesian_data]
            magnitudes = [np.abs(coord[2]) for coord in original_cartesian_data]  # 使用复数的幅度

            plt.figure(figsize=(8, 8))
            scatter = plt.scatter(x_values, y_values, c=magnitudes, cmap='viridis', marker='o')
            plt.title("Cartesian Coordinates (Original Data)")
            plt.xlabel("X")
            plt.ylabel("Y")
            plt.axis("equal")
            plt.colorbar(scatter, label="Magnitude")
            plt.show()


def main():
    # 创建多个XRayScan对象，增加点的数量
    angles = np.linspace(0, 180, num=12)  # 12个不同角度
    scans = []
    for angle in angles:
        sensor_data = [0, 1, 0, 3, 0, 2, 0, 1, 0, 4, 0, 1, 0, 2, 0, 3, 0, 1, 0]  # 增加数据点
        scans.append(XRayScan(angle, sensor_data))

    # 创建XRayScanCollection对象并添加扫描
    collection = XRayScanCollection()
    for scan in scans:
        collection.add_scan(scan)

    # 创建XRayScanCollectionFFT对象，并计算傅里叶变换
    fft_collection = XRayScanCollectionFFT(collection)
    fft_collection.compute_fft()

    # 绘制位于指定长方形内的点和长方形外框
    aspect_ratio = 1.0  # 1:1 长宽比
    fft_collection.plot_filtered_cartesian(aspect_ratio)

    # 获取存储在矩形区域内的点，并创建RectangleData对象
    rectangle_data = fft_collection.get_rectangle_data(aspect_ratio)
    print(rectangle_data)

    # 打印存储的点
    for point in rectangle_data.get_points():
        print(point)

if __name__ == "__main__":
    main()